# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""The module is used to register MindSpore builtin data augment APIs."""

import inspect
import random
from re import S
from PIL.Image import BILINEAR, new
import cv2
import numpy as np
from mindspore import Tensor
import numbers

import mindspore.dataset.vision.c_transforms as C
import mindspore.dataset.vision.py_transforms as PY
import mindspore.dataset.transforms.py_transforms as PY_trans
from numpy.lib.arraysetops import isin

from mindvideo.common.utils.class_factory import ClassFactory, ModuleType


def register_builtin_transforms():
    """ register MindSpore builtin dataset class. """
    for module_name in set(dir(C) + dir(PY)):
        if module_name.startswith('__'):
            continue
        transforms = getattr(C, module_name, None) if getattr(C, module_name, None) else getattr(PY, module_name)
        if inspect.isclass(transforms):
            ClassFactory.register_cls(transforms, ModuleType.PIPELINE)

@ClassFactory.register(ModuleType.PIPELINE)
class ReOrder(PY_trans.PyTensorOperation):
    """Rearrange the order of dims of data.
    Args:
        new_order(tuple), new_order of output.
    """
    def __init__(self, order):
        self.order = tuple(order)
    def __call__(self, x):
        if isinstance(x, np.ndarray):
            return np.transpose(x, self.order)
        elif isinstance(x, Tensor):
            return x.transpose(x, self.order)
        else:
            raise AssertionError("""The type of input should be np.array or 
                                    mindspore.Tensor but got {}."""\
                                    .format(type(x).__name__))   
@ClassFactory.register(ModuleType.PIPELINE)
class VideoToTensor(PY_trans.PyTensorOperation):
    """Convert the input video frames in type numpy.ndarray of shape (T, H, W, C)
       in the range [0, 255] to numpy.ndarray of shape (C, T, H, W)  in the range
       [0.0, 1.0] with the desired dtype.
    Args:
        new_order(tuple), new_order of output.
    """
    # def __init__(self, order):
    #     self.order = tuple(order)
    def __call__(self, x):
        if isinstance(x, np.ndarray):
            x = x/255.0
            return np.transpose(x, (3, 0, 1, 2)).astype(np.float32)
        else:
            raise AssertionError("""The type of input should be numpy.ndarray 
                                    but got {}.""".format(type(x).__name__))

@ClassFactory.register(ModuleType.PIPELINE)
class VideoResize(PY_trans.PyTensorOperation):
    """Resize the given video sequences (t x h x w x c) at the given size.
    Args:
        size(tuple or int): Desired output size after resize. 
    """

    def __init__(self, size, interpolation="bilinear"):
        self.size = size
        self.inter = interpolation
        self._opencv_mindspore_interpolation_map = {
            "nearest": cv2.INTER_NEAREST,
            "linear": cv2.INTER_LINEAR,
            "bilinear": cv2.INTER_AREA,
            "bicubic": cv2.INTER_CUBIC,
        }
    def __call__(self, x: np.ndarray) -> np.ndarray:
        resized_img_array_list = [
            cv2.resize(
                img_array,
                self.size,  # The input order for OpenCV is w, h.
                interpolation=self._opencv_mindspore_interpolation_map[self.inter],
            ) 
            for img_array in x
        ]
        
        img_array = np.concatenate(
            [np.expand_dims(img_array, axis=0) for img_array in resized_img_array_list],
            axis=0,
        )
        return img_array

@ClassFactory.register(ModuleType.PIPELINE)
class VideoShortEdgeResize(PY_trans.PyTensorOperation):
    """Resize the given video sequences (t x h x w x c) at the given size.
       And make sure the smallest dimension in (h, w) is 256 pixels.
    Args:
        size(int): Desired output size after resize. 
    """
    def __init__(self, size, interpolation="bilinear"):
        self.size = size
        self.inter = interpolation
        self._opencv_mindspore_interpolation_map = {
            "nearest": cv2.INTER_NEAREST,
            "linear": cv2.INTER_LINEAR,
            "bilinear": cv2.INTER_AREA,
            "bicubic": cv2.INTER_CUBIC,
        }
    def __call__(self, x: np.ndarray) -> np.ndarray:
        _, h, w, _ = x.shape
        if h<w:
            scale = self.size*1.0/h
        else:
            scale = self.size*1.0/w
        new_size = (int(scale*h), int(scale*w))    
        # print(self.size)    
        resized_img_array_list = [
            cv2.resize(
                img_array,
                new_size,  # The input order for OpenCV is w, h.
                interpolation=self._opencv_mindspore_interpolation_map[self.inter],
            ) 
            for img_array in x
        ]
        
        img_array = np.concatenate(
            [np.expand_dims(img_array, axis=0) for img_array in resized_img_array_list],
            axis=0,
        )
        return img_array
        
@ClassFactory.register(ModuleType.PIPELINE)
class TensorResize(PY_trans.PyTensorOperation):
    """Resize the seq images.
    Args: 
        size(tuple or int): Desired output size after resize. 
    """
    
    def __init__(self, size):
        self.size = size
    def __call__(self, x: Tensor) -> Tensor:
        #mindspore.Tensor.resize  input&output(tensor)
        x = x.resize(self.size)
        return x

@ClassFactory.register(ModuleType.PIPELINE)
class VideoRandomCrop(PY_trans.PyTensorOperation):
    """Crop the given video sequences (c x t x h x w) at a random location.
    Args:
        size (sequence or int): Desired output size of the crop. If size is an
            int instead of sequence like (h, w), a square crop (size, size) is
            made.
    """

    def __init__(self, size):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = tuple(size)

    @staticmethod
    def get_params(img, output_size):
        """Get parameters for ``crop`` for a random crop.
        Args:
            img (PIL Image): Image to be cropped.
            output_size (tuple): Expected output size of the crop.
        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
        """
        t, h, w, c = img.shape
        th, tw = output_size
        if w == tw and h == th:
            return 0, 0, h, w

        i = random.randint(0, h - th) if h!=th else 0
        j = random.randint(0, w - tw) if w!=tw else 0
        return i, j, th, tw

    def __call__(self, imgs):
        
        i, j, h, w = self.get_params(imgs, self.size)

        imgs = imgs[:, i:i+h, j:j+w, :]
        return imgs

    def __repr__(self):
        return self.__class__.__name__ + '(size={0})'.format(self.size)

@ClassFactory.register(ModuleType.PIPELINE)
class VideoCentorCrop(PY_trans.PyTensorOperation):
    """Crops the given seq Images at the center.
    Args:
        size (sequence or int): Desired output size of the crop. 
    """

    def __init__(self, size = (224, 224)):
        self.size = size

    def __call__(self, imgs):
        """
        Args:
            img (PIL Image): Image to be cropped.
        Returns:
            PIL Image: Cropped image.
        """
        t, h, w, c = imgs.shape
        th, tw = self.size
        i = int(np.round((h - th) / 2.))
        j = int(np.round((w - tw) / 2.))

        return imgs[:, i:i+th, j:j+tw, :]

@ClassFactory.register(ModuleType.PIPELINE)
class VideoRandomHorizontalFlip(PY_trans.PyTensorOperation):
    """Horizontally flip the given seq Images randomly with a given probability.
    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, prob=0.5):
        self.p = prob

    def __call__(self, imgs):
        """
        Args:
            img (seq Images): seq Images to be flipped.
        Returns:
            seq Images: Randomly flipped seq images.
        """
        if random.random() < self.p:
            # t x h x w
            return np.flip(imgs, axis=2).copy()
        return imgs
